from datasets import Dataset
from transformers import AutoTokenizer, AutoModelForCausalLM, DataCollatorForSeq2Seq, TrainingArguments, Trainer
from peft import LoraConfig, TaskType, get_peft_model
import pandas as pd
import torch
import os

os.environ['WANDB_DISABLED'] = 'True'


# 数据处理函数
def process_func(example):
    MAX_LENGTH = 512
    instruction = tokenizer.encode(text="\n".join(["<|system|>", "You are now a financial expert, trying to help me answer my financial questions", "<|user|>",
                                                   example["instruction"] + example["input"] + "<|assistant|>"]).strip() + "\n",
                                   add_special_tokens=True, truncation=True, max_length=MAX_LENGTH)
    response = tokenizer.encode(text=example["output"], add_special_tokens=False, truncation=True, max_length=MAX_LENGTH)
    input_ids = instruction + response + [tokenizer.eos_token_id]
    labels = [tokenizer.pad_token_id] * len(instruction) + response + [tokenizer.eos_token_id]
    pad_len = MAX_LENGTH - len(input_ids)
    # print()
    input_ids += [tokenizer.pad_token_id] * pad_len
    labels += [tokenizer.pad_token_id] * pad_len
    labels = [(l if l != tokenizer.pad_token_id else -100) for l in labels]
    return {
        "input_ids": input_ids,
        "labels": labels
    }


# LoraConfig
config = LoraConfig(task_type=TaskType.CAUSAL_LM, target_modules=["query_key_value","lm_head"],modules_to_save=["word_embeddings"])


# 配置训练参数
args = TrainingArguments(
    output_dir="output",
    per_device_train_batch_size=1,
    gradient_accumulation_steps=8,
    logging_steps=20,
    num_train_epochs=1
)

if __name__ == '__main__':

    # 数据和模型地址
    data_path = 'data/split_data.json'
    MODEL_PATH = 'F:\pretrain_model\ZhipuAI\chatglm3-6b'

    # 处理数据集
    df = pd.read_json(data_path)
    ds = Dataset.from_pandas(df)

    # 加载模型和tokenizer
    model = AutoModelForCausalLM.from_pretrained(pretrained_model_name_or_path=MODEL_PATH, trust_remote_code=True,
                                                 device_map='auto',low_cpu_mem_usage=True,torch_dtype=torch.bfloat16)
    tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name_or_path=MODEL_PATH, trust_remote_code=True,
                                              use_fast=False)
    # tokenizer.pad_token_id = tokenizer.eod_id

    tokenized_id = ds.map(process_func,remove_columns=ds.column_names)

    #model.enable_input_require_grads()
    data_collator = DataCollatorForSeq2Seq(
        tokenizer,
        model=model,
        label_pad_token_id=-100,
        pad_to_multiple_of=None,
        padding=False
    )

    model = get_peft_model(model, config)

    trainer = Trainer(
        model=model,
        args=args,
        train_dataset=tokenized_id,
        data_collator=data_collator,
    )
    trainer.train()  # 开始训练
    # response, history = model.chat(tokenizer, "你是谁", history=[], system="现在你要扮演皇帝身边的女人--甄嬛.")
    # print(response)

